#!/bin/bash

# 工程实践平台Docker镜像下载工具
# 使用方法: ./scripts/download-images.sh

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查是否为root用户
check_root() {
    if [[ $EUID -eq 0 ]]; then
        log_error "请不要使用root用户运行此脚本"
        exit 1
    fi
}

# 检查Docker服务状态
check_docker() {
    log_info "检查Docker服务状态..."
    
    if ! docker info &> /dev/null; then
        log_error "Docker服务未启动"
        echo "请启动Docker服务:"
        echo "  Linux: sudo systemctl start docker"
        echo "  Windows/macOS: 启动Docker Desktop"
        exit 1
    fi
    
    log_success "Docker服务运行正常"
}

# 配置Docker镜像加速器
configure_docker_mirrors() {
    log_info "配置Docker镜像加速器..."
    
    if [[ ! -f /etc/docker/daemon.json ]]; then
        sudo mkdir -p /etc/docker
        
        cat > /tmp/daemon.json << 'EOF'
{
  "registry-mirrors": [
    "https://docker.mirrors.ustc.edu.cn",
    "https://hub-mirror.c.163.com",
    "https://mirror.baidubce.com",
    "https://registry.docker-cn.com"
  ],
  "log-driver": "json-file",
  "log-opts": {
    "max-size": "100m",
    "max-file": "3"
  }
}
EOF
        
        sudo mv /tmp/daemon.json /etc/docker/daemon.json
        
        log_info "重启Docker服务以应用配置..."
        sudo systemctl restart docker
        sleep 10
        
        log_success "Docker镜像加速器配置完成"
    else
        log_info "Docker镜像加速器已配置"
    fi
}

# 定义镜像列表
define_images() {
    # 核心服务镜像
    CORE_IMAGES=(
        "mysql:5.7"
        "redis:7-alpine"
        "rabbitmq:3-management"
        "elasticsearch:7.17.0"
    )
    
    # 开发环境镜像
    DEV_IMAGES=(
        "codercom/code-server:latest"
        "jupyter/datascience-notebook:latest"
        "dorowu/ubuntu-desktop-lxde-vnc:latest"
    )
    
    # 监控镜像
    MONITORING_IMAGES=(
        "prom/prometheus:latest"
        "grafana/grafana:latest"
    )
    
    # 合并所有镜像
    ALL_IMAGES=("${CORE_IMAGES[@]}" "${DEV_IMAGES[@]}" "${MONITORING_IMAGES[@]}")
    
    echo "${ALL_IMAGES[@]}"
}

# 创建镜像目录
create_image_directories() {
    log_info "创建镜像目录..."
    
    mkdir -p images/{core,dev,monitoring}
    
    log_success "镜像目录创建完成"
}

# 下载镜像
download_image() {
    local image=$1
    local category=$2
    local max_attempts=3
    
    log_info "下载镜像: $image"
    
    for attempt in {1..max_attempts}; do
        if docker pull "$image"; then
            log_success "成功下载: $image"
            
            # 保存镜像到文件
            local filename=$(echo "$image" | sed 's/[:\/]/_/g').tar
            local filepath="images/$category/$filename"
            
            log_info "保存镜像到: $filepath"
            if docker save "$image" > "$filepath"; then
                log_success "成功保存: $filename"
                
                # 显示文件大小
                local size=$(du -h "$filepath" | cut -f1)
                echo "   文件大小: $size"
                
                return 0
            else
                log_error "保存失败: $filename"
                return 1
            fi
        else
            log_warning "尝试 $attempt/$max_attempts 失败"
            if [[ $attempt -lt $max_attempts ]]; then
                echo "等待重试..."
                sleep 10
            else
                log_error "无法下载镜像: $image"
                return 1
            fi
        fi
    done
}

# 批量下载镜像
download_all_images() {
    log_info "开始批量下载镜像..."
    
    local images=($(define_images))
    local total=${#images[@]}
    local success=0
    local failed=0
    
    echo "总共需要下载 $total 个镜像"
    echo
    
    # 下载核心服务镜像
    log_info "下载核心服务镜像..."
    for image in "mysql:5.7" "redis:7-alpine" "rabbitmq:3-management" "elasticsearch:7.17.0"; do
        if download_image "$image" "core"; then
            ((success++))
        else
            ((failed++))
        fi
        echo
    done
    
    # 下载开发环境镜像
    log_info "下载开发环境镜像..."
    for image in "codercom/code-server:latest" "jupyter/datascience-notebook:latest" "dorowu/ubuntu-desktop-lxde-vnc:latest"; do
        if download_image "$image" "dev"; then
            ((success++))
        else
            ((failed++))
        fi
        echo
    done
    
    # 下载监控镜像
    log_info "下载监控镜像..."
    for image in "prom/prometheus:latest" "grafana/grafana:latest"; do
        if download_image "$image" "monitoring"; then
            ((success++))
        else
            ((failed++))
        fi
        echo
    done
    
    # 显示下载结果
    echo "=========================================="
    echo "下载完成！"
    echo "=========================================="
    echo "成功: $success"
    echo "失败: $failed"
    echo "总计: $total"
    echo "=========================================="
}

# 显示镜像文件信息
show_image_files() {
    log_info "镜像文件信息..."
    
    echo
    echo "=========================================="
    echo "镜像文件列表"
    echo "=========================================="
    
    local total_size=0
    local total_files=0
    
    for category in core dev monitoring; do
        if [[ -d "images/$category" ]]; then
            echo
            echo "$(echo $category | tr '[:lower:]' '[:upper:]') 镜像:"
            if [[ -n "$(ls -A images/$category)" ]]; then
                for file in images/$category/*.tar; do
                    if [[ -f "$file" ]]; then
                        local size=$(du -h "$file" | cut -f1)
                        local filename=$(basename "$file")
                        echo "  $filename ($size)"
                        ((total_files++))
                        # 累加文件大小（以KB为单位）
                        total_size=$((total_size + $(du -k "$file" | cut -f1)))
                    fi
                done
            else
                echo "  无镜像文件"
            fi
        fi
    done
    
    echo
    echo "总计: $total_files 个文件"
    echo "总大小: $(numfmt --to=iec $((total_size * 1024)))"
    echo "=========================================="
}

# 显示使用方法
show_usage() {
    echo
    echo "=========================================="
    echo "使用方法"
    echo "=========================================="
    echo "1. 离线部署:"
    echo "   - 将整个项目目录复制到目标服务器"
    echo "   - 运行: ./scripts/offline-deploy.sh"
    echo
    echo "2. 手动加载镜像:"
    echo "   # 加载所有镜像"
    echo "   for file in images/*/*.tar; do docker load < \$file; done"
    echo
    echo "   # 按类别加载"
    echo "   for file in images/core/*.tar; do docker load < \$file; done"
    echo "   for file in images/dev/*.tar; do docker load < \$file; done"
    echo "   for file in images/monitoring/*.tar; do docker load < \$file; done"
    echo
    echo "3. 验证镜像:"
    echo "   docker images | grep -E '(mysql|redis|rabbitmq|elasticsearch|code-server|jupyter|ubuntu-vnc|prometheus|grafana)'"
    echo "=========================================="
}

# 清理临时文件
cleanup() {
    log_info "清理临时文件..."
    
    # 清理可能存在的临时文件
    rm -f /tmp/daemon.json
    
    log_success "清理完成"
}

# 主函数
main() {
    echo "=========================================="
    echo "工程实践平台Docker镜像下载工具"
    echo "=========================================="
    echo
    
    # 检查环境
    check_root
    check_docker
    
    # 配置镜像加速器
    configure_docker_mirrors
    
    # 创建目录
    create_image_directories
    
    # 下载镜像
    download_all_images
    
    # 显示结果
    show_image_files
    show_usage
    
    # 清理
    cleanup
    
    log_success "镜像下载工具执行完成！"
}

# 执行主函数
main "$@" 